package main

import (
	"crypto/rand"
	"encoding/csv"
	"fmt"
	"io"
	"log"
	"os"
	"pvpkeet-go/pvpkeet" // 确保这里的模块名与你的 go.mod 文件一致
	"time"
)

// PerformanceResult (已更新)
type PerformanceResult struct {
	Experiment     string
	Parameter      string
	Value          int
	ProofGenTimeMs float64
	VerifyTimeMs   float64
	// SetupTimeMs       float64
	// EncryptTimeMs     float64
	// AvgTestTimeMs     float64
	// TotalTestTimeMs   float64
	// AvgProofTimeMs    float64
	// TotalProofTimeMs  float64
	// AvgVerifyTimeMs   float64
	// TotalVerifyTimeMs float64
}

// 辅助函数：生成随机数据
func generateDummyFile(size int) []byte {
	data := make([]byte, size)
	rand.Read(data)
	return data
}

// 实验1：考察文件大小的影响
func fileSizeExp(writer *csv.Writer, protocol *pvpkeet.PVPKEET) {
	log.Println("\n\n<<<<<<<<<< PVPKEET Experiment 1: Performance vs. File Size >>>>>>>>>>")
	const K = 5
	fileSizeOptionsMB := []int{4, 8, 12, 16, 20}

	for _, sizeMB := range fileSizeOptionsMB {
		fileSize := sizeMB * 1024 * 1024
		log.Printf("\n--- Testing with FileSize: %dMB ---", sizeMB)

		dummyFile := make([]byte, fileSize)
		rand.Read(dummyFile)

		userKeys := protocol.UserKeyGen()
		cloudKeys := make([]*pvpkeet.CloudKeyPair, K)
		for i := 0; i < K; i++ {
			cloudKeys[i] = protocol.CloudKeyGen()
		}

		ciphertexts := make([]*pvpkeet.Ciphertext, K)
		for i := 0; i < K; i++ {
			ciphertexts[i] = protocol.Encryption(userKeys.Pku, cloudKeys[i].Pkc, dummyFile)
		}

		ctA, ctB := ciphertexts[0], ciphertexts[1]
		skcA, skcB := cloudKeys[0].Skc, cloudKeys[1].Skc
		pkcA := cloudKeys[0].Pkc

		protocol.Test(ctA, ctB, skcA, skcB)

		startProof := time.Now()
		var proof *pvpkeet.Proof
		// TODO:证明k-1次
		for range K - 1 {
			proof = protocol.Proof(ctA, ctB, skcA, skcB, userKeys.Pku)
		}
		proofGenTime := time.Since(startProof)

		startVerify := time.Now()
		// TODO:验证k-1次
		for range K - 1 {
			protocol.Verify(proof, ctA, ctB, pkcA, userKeys.Pku)
		}
		verifyTime := time.Since(startVerify)

		result := PerformanceResult{
			Experiment: "FileSize", Parameter: "FileSize(MB)", Value: sizeMB,
			ProofGenTimeMs: float64(proofGenTime.Microseconds()) / 1000.0,
			VerifyTimeMs:   float64(verifyTime.Microseconds()) / 1000.0,
			// SetupTimeMs:       float64(setupTime.Microseconds()) / 1000.0,
			// EncryptTimeMs:     float64(avgEncryptTime.Microseconds()) / 1000.0,
			// AvgTestTimeMs:     float64(testTime.Microseconds()) / 1000.0,
			// TotalTestTimeMs:   float64(testTime.Microseconds()) / 1000.0,
			// AvgProofTimeMs:    float64(proofTime.Microseconds()) / 1000.0,
			// TotalProofTimeMs:  float64(proofTime.Microseconds()) / 1000.0,
			// AvgVerifyTimeMs:   float64(verifyTime.Microseconds()) / 1000.0,
			// TotalVerifyTimeMs: float64(verifyTime.Microseconds()) / 1000.0,
		}
		log.Printf("[PERF] Result: %+v", result)
		writer.Write([]string{
			result.Experiment, result.Parameter, fmt.Sprintf("%d", result.Value),
			fmt.Sprintf("%.3f", result.ProofGenTimeMs),
			fmt.Sprintf("%.3f", result.VerifyTimeMs),
			// fmt.Sprintf("%.3f", result.SetupTimeMs), fmt.Sprintf("%.3f", result.EncryptTimeMs),
			// fmt.Sprintf("%.3f", result.AvgTestTimeMs), fmt.Sprintf("%.3f", result.TotalTestTimeMs),
			// fmt.Sprintf("%.3f", result.AvgProofTimeMs), fmt.Sprintf("%.3f", result.TotalProofTimeMs),
			// fmt.Sprintf("%.3f", result.AvgVerifyTimeMs), fmt.Sprintf("%.3f", result.TotalVerifyTimeMs),
		})
		writer.Flush()
	}
}

// 实验2：考察 CSP 数量的影响
func cspCountExp(writer *csv.Writer, protocol *pvpkeet.PVPKEET) {
	log.Println("\n\n<<<<<<<<<< PVPKEET Experiment 2: Performance vs. Number of CSPs >>>>>>>>>>")
	const FileSize = 4 * 1024 * 1024
	cspCounts := []int{5, 10, 15, 20, 25, 30}

	dummyFile := make([]byte, FileSize)
	rand.Read(dummyFile)

	for _, k := range cspCounts {
		log.Printf("\n--- Testing with %d CSPs ---", k)

		userKeys := protocol.UserKeyGen()
		cloudKeys := make([]*pvpkeet.CloudKeyPair, k)
		for i := 0; i < k; i++ {
			cloudKeys[i] = protocol.CloudKeyGen()
		}

		ciphertexts := make([]*pvpkeet.Ciphertext, k)
		for i := 0; i < k; i++ {
			ciphertexts[i] = protocol.Encryption(userKeys.Pku, cloudKeys[i].Pkc, dummyFile)
		}

		ctA, ctB := ciphertexts[0], ciphertexts[1]
		skcA, skcB := cloudKeys[0].Skc, cloudKeys[1].Skc
		pkcA := cloudKeys[0].Pkc

		protocol.Test(ctA, ctB, skcA, skcB)

		startProof := time.Now()
		var proof *pvpkeet.Proof
		// TODO:证明k-1次
		for range k - 1 {
			proof = protocol.Proof(ctA, ctB, skcA, skcB, userKeys.Pku)
		}
		proofGenTime := time.Since(startProof)

		startVerify := time.Now()
		// TODO:验证k-1次
		for range k - 1 {
			protocol.Verify(proof, ctA, ctB, pkcA, userKeys.Pku)
		}
		verifyTime := time.Since(startVerify)

		result := PerformanceResult{
			Experiment: "CSPCount", Parameter: "NumCSPs", Value: k,
			ProofGenTimeMs: float64(proofGenTime.Microseconds()) / 1000.0,
			VerifyTimeMs:   float64(verifyTime.Microseconds()) / 1000.0,
		}
		log.Printf("[PERF] Result: %+v", result)
		writer.Write([]string{
			result.Experiment, result.Parameter, fmt.Sprintf("%d", result.Value),
			fmt.Sprintf("%.3f", result.ProofGenTimeMs),
			fmt.Sprintf("%.3f", result.VerifyTimeMs),
		})
		writer.Flush()
	}
}

func main() {
	logFile, err := os.OpenFile("experiment_pvpkeet.log", os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
	if err != nil {
		log.Fatalf("Failed to open log file: %v", err)
	}
	defer logFile.Close()
	mw := io.MultiWriter(os.Stdout, logFile)
	log.SetOutput(mw)

	csvFile, err := os.Create("results_pvpkeet.csv")
	if err != nil {
		log.Fatalf("Failed to create CSV file: %v", err)
	}
	defer csvFile.Close()
	csvWriter := csv.NewWriter(csvFile)
	defer csvWriter.Flush()
	csvWriter.Write([]string{"Experiment", "Parameter", "Value", "ProofGenTime(ms)", "VerifyTime(ms)"})

	protocol := pvpkeet.Setup("SS512")

	fileSizeExp(csvWriter, protocol)
	cspCountExp(csvWriter, protocol)

	log.Println("\n\n<<<<<<<<<< All PVPKEET experiments finished. >>>>>>>>>>")
}
